/**
 * @file ZSuffixTree.h
 * @brief suffix tree main code
 * @author babear.wangz
 8 @date 2011-04-08
 */

#ifndef ZST_ZSuffixTree_H
#define ZST_ZSuffixTree_H

#include "ZTraits.h"
#include "ZSTEdgeAndNode.h"
#include "ZSTSequenceBuffer.h"
#include "ZSTEdgeBuffer.h"

/**
 * @class ZSuffixTree
 * @brief suffix tree
 * @author babear.wangz
 * @date 2011-04-08
 */
template <typename SequenceElementType, // ordinarily char
    typename SpaceParameter/* = ZSTBufferSpaceParameter */,
    typename SequenceClass/* = ZSTSequenceBuffer<SequenceElementType,
        ZTypeTraits<SequenceElementType>, SpaceParameter> */,
    uint32_t WordLengthFunc(SequenceClass*, size_t), // word seg method
    typename SequenceHashMethod = SequenceHash<SequenceElementType> >
class ZSuffixTree {
public:
    // default constructor
    ZSuffixTree()
    : edges(&sequence, &nodes)
#ifdef DEBUG
    , progress(0)
#endif
    {
        Init();
    }
    /**
     * @brief constructor
     * @param seq the sequence to be copied to suffix tree
     * @param l   the length of the sequence
     */
    ZSuffixTree(const SequenceElementType* seq, uint32_t l)
    : sequence(seq, l)
    , edges(&sequence, &nodes)
#ifdef DEBUG
    , progress(0)
#endif
    {
        Init();
    }
public:
    /**
     * @function Construct
     * @brief construct suffix tree
     * @author babear.wangz
     * @date 2011-04-13
     */
    bool Construct();
    /**
     * @function PushBack
     * @brief push back a new element and update the suffix tree
     * @param elem pointer to the new elment
     * @param l    the length of the element sequence.
     * @author babear.wangz
     */
    bool PushBack(const SequenceElementType* elem, uint32_t l);
    /**
     * @function ReConstruct
     * @brief rebuild a suffix tree with a new sequence
     * @param seq the new sequence
     * @param l the length of new sequence
     * @author babear.wangz
     * @date 2011-04-13
     */
    bool ReConstruct(const SequenceElementType* seq, uint32_t l);

    /**
     * @function 
     * @author babear.wangz
     */

    /**
     * @function Swap
     * @brief swap two suffix trees
     * @author babear.wangz
     */
    void Swap(ZSuffixTree& suftree)
    {
        sequence.Swap(suftree.sequence);
        nodes.Swap(suftree.nodes);
        edges.SwapWithNoSequenceAndNodes(suftree.edges);
        ::Swap(active.node, suftree.active.node);
        ::Swap(active.edge, suftree.active.edge);
        ::Swap(active.len, suftree.active.len);
#ifdef DEBUG
        ::Swap(progress, suftree.progress);
#endif
    }

private:
    /**
     * @brief update the suffix tree through accessing seq_idx element
     * @author babear.wangz
     */
    bool Update(uint32_t seq_idx);
    /**
     * @brief check if the seq_idx should split the tree in current station
     * @author babear.wangz
     */
    bool ShouldSplit(uint32_t seq_idx, uint32_t wlen);
    /**
     * @brief initialize the active and root. called by constructors
     * @author babear.wangz
     */
    bool Init();
    /**
     * @brief canonize the inner point representation
     * @author babear.wangz
     */
    void Canonize(uint32_t& nid, uint32_t& s, uint32_t& e, uint32_t& eid);

    /**
     * @fn CheckEdgeAdvance
     * @brief check if the edge should be advanced if the input is sidx
     * @brief Check* mainly used in ASSERT
     * @param eid the edge index
     * @param len the current inner point on the edge
     * @param sidx the next input
     * @param slen the word length
     * @author babear.wangz
     * @date 2011-04-12
     */
    bool CheckEdgeAdvance(uint32_t eid,
                          uint32_t len,
                          uint32_t sidx,
                          uint32_t slen);
private:
    DISALLOW_COPY_AND_ASSIGN(ZSuffixTree);
    // the sequence
    SequenceClass sequence;
    // the nodes
    ZSTSequenceBuffer<ZSTNode<uint32_t> > nodes;
    // the edges
    ZSTEdgeBuffer<uint32_t, NodeIDHash<uint32_t>,
        SequenceElementType,
        SequenceClass,
        SequenceHashMethod,
        WordLengthFunc,
        SpaceParameter> edges;
    // the active point
    struct Active {
        uint32_t node;
        uint32_t edge;
        uint32_t len;
    } active;
#ifdef DEBUG
    // build progress
    uint32_t progress;
#endif
};

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>::Construct()
{
#ifdef DEBUG
    if (progress != 0) CRASH();
#endif
    for (uint32_t i = 0; i < sequence.Size(); ) {
        if (!Update(i)) {
#ifdef DEBUG
            progress = i;
#endif
            return false;
        }
        i += WordLengthFunc(&sequence, i);
    }
#ifdef DEBUG
    progress = sequence.Size();
#endif
    return true;
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::PushBack(const SequenceElementType* elem, uint32_t l)
{
    if (progress < sequence.Size()) CRASH();
    uint32_t idx = sequence.Size();
    for (uint32_t i = 0; i < l; ++i) {
        sequence.PushBack(elem[i]);
    }
    for (uint32_t i = idx; i < sequence.Size(); ) {
        if (!Update(i)) {
            return false;
        }
        i += WordLengthFunc(&sequence, i);
    }
    progress = sequence.Size(); 
    return true;
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::ReConstruct(const SequenceElementType* seq, uint32_t l)
{
    {
        ZSuffixTree newst(seq, l);
        this->Swap(newst);
    }
    return Construct();
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::Update(uint32_t seq_idx)
{
    uint32_t wlen = WordLengthFunc(&sequence, seq_idx);
    bool should_split = ShouldSplit(seq_idx, wlen);//meanwhile, update active
    if (!should_split) return true;
    uint32_t t_nid = active.node;
    uint32_t t_s, t_e, t_ed;
    if (active.len == 0) {
        t_s = t_e = -1;
    } else {
        ZSTEdge<uint32_t>& ed = edges.GetEdge(active.edge);
        t_s = ed.start;
        ASSERT(ed.start + active.len - 1 < ed.end);//first split is canonized
        t_e = ed.start + active.len - 1;
    }
    t_ed = active.edge;
    uint32_t pre_nid = -1;
    while (1) {
        if (t_s == (uint32_t)(-1) && t_e == (uint32_t)(-1)) {
            // directly splitting
            edges.Insert(t_nid, seq_idx, -1, -1);
            if (pre_nid != (uint32_t)(-1)) {
                nodes[pre_nid].suflink = t_nid;
            }
            pre_nid = t_nid;
            t_nid = nodes[t_nid].suflink;
        } else {
            uint32_t new_nid = nodes.Size();
            nodes.PushBack(ZSTNode<uint32_t>());
            uint32_t t_enode, t_end, t_splitstart;
            {//ed must be used before Insert method, because insert method
             //maybe change the ed reference
                ZSTEdge<uint32_t>& ed = edges.GetEdge(t_ed);
                t_enode = ed.enode;
                t_end = ed.end;
                ed.enode = new_nid;
                ed.end = ed.start + t_e - t_s;
                t_splitstart = ed.end + 1;
                ASSERT(t_end == (uint32_t)(-1) || t_end > ed.end);
            }
            edges.Insert(new_nid, seq_idx, -1, -1);
            edges.Insert(new_nid, t_splitstart, t_end, t_enode);
            // update suffix link
            if (pre_nid != (uint32_t)(-1)) {
                nodes[pre_nid].suflink = new_nid;
            }
            // suffix link
            pre_nid = new_nid;
            t_nid = nodes[t_nid].suflink;
        }
        if (t_nid == (uint32_t)(-1)) {//last one is root
            if (t_s == (uint32_t)(-1) && t_e == (uint32_t)(-1)) {
                active.node = 0;
                active.edge = 0;
                active.len = 0;
                break;
            } else {
                // shrink [t_s, t_e]
                ASSERT(t_s <= t_e);
                uint32_t t_wlen = WordLengthFunc(&sequence, t_s);
                ASSERT(t_s + t_wlen - 1 <= t_e);
                t_s += t_wlen;
                t_nid = 0;
                if (t_s > t_e) {
                    t_s = t_e = -1;
                } else {
                    Canonize(t_nid, t_s, t_e, t_ed);
                }
            }
        } else {
            Canonize(t_nid, t_s, t_e, t_ed);
        }
        // judge if it should continue to split
        if (t_s == (uint32_t)(-1) && t_e == (uint32_t)(-1)) {
            uint32_t t_eid = -1;
            ZSTEdge<uint32_t>& ed = edges.Find(t_nid, seq_idx, t_eid);
            if (!ed.IsEmpty()) {
                // if the edge is existed, and this update should be stopped.
                // and the t_nid should be the suflink of pre_nid
                if (pre_nid != (uint32_t)(-1)) {
                    nodes[pre_nid].suflink = t_nid;
                }
                t_s = seq_idx;
                t_e = seq_idx + wlen - 1;
                Canonize(t_nid, t_s, t_e, t_eid);
                active.node = t_nid;
                if (t_s == (uint32_t)(-1) && t_e == (uint32_t)(-1)) {
                    active.edge = active.len = 0;
                } else {
                    active.edge = t_eid;
                    active.len = t_e - t_s + 1;
                }
                break;
            }
        } else {
            // it is impossible that previous node splited but current is
            // a edge and can be advanced.
            ASSERT(!CheckEdgeAdvance(t_ed, t_e - t_s + 1, seq_idx, wlen));
        }
    }
    return true;
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::CheckEdgeAdvance(uint32_t eid, uint32_t len,
                       uint32_t sidx, uint32_t slen)
{
    ZSTEdge<uint32_t>& e = edges.GetEdge(eid);
    ASSERT(e.start + len <= e.end);// must be a inner point
    if (e.end - (e.start + len) + 1 < slen) return false;
    if (!sequence.Equal(e.start + len, sidx, slen)) return false;
    return true;
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
void ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::Canonize(uint32_t& nid, uint32_t& s, uint32_t& e, uint32_t& eid)
{
    if (s == (uint32_t)(-1) && e == (uint32_t)(-1)) return;
    while (1) {
        uint32_t wlen = WordLengthFunc(&sequence, s);
        ZSTEdge<uint32_t>& ed = edges.Find(nid, s, eid);
        ASSERT(!ed.IsEmpty());
        ASSERT(ed.enode == (uint32_t)(-1) && ed.end == (uint32_t)(-1) ||
               ed.enode != (uint32_t)(-1) && ed.end != (uint32_t)(-1));
        if (ed.enode == (uint32_t)(-1)) return;
        if (e - s < ed.end - ed.start) {
            break;
        } else if (e - s == ed.end - ed.start) {// pure node
            nid = ed.enode;
            s = e = eid = -1;
            break;
        } else {
            nid = ed.enode;
            s += ed.end - ed.start + 1;
        }
    }
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>
    ::ShouldSplit(uint32_t seq_idx, uint32_t wlen)
{
    if (active.len == 0) {
        uint32_t eid = -1;
        ZSTEdge<uint32_t>& e = edges.Find(active.node, seq_idx, eid);
        if (!e.IsEmpty()) {
            if (e.end == (uint32_t)(-1)) {
                ASSERT(e.enode == (uint32_t)(-1));
                active.edge = eid;
                active.len = wlen;
            } else {
                ASSERT(e.enode != (uint32_t)(-1));
                uint32_t elen = e.end - e.start + 1;
                ASSERT(elen >= wlen);
                if (elen == wlen) {
                    active.node = e.enode;
                    active.edge = 0;
                    active.len = 0;
                } else {
                    active.edge = eid;
                    active.len = wlen;
                }
            }
            return false;
        } else {
            return true;
        }
    } else {
        ZSTEdge<uint32_t>& e = edges.GetEdge(active.edge);
        ASSERT(e.node == active.node);
        ASSERT(!e.IsEmpty());
        if (e.end == (uint32_t)(-1)) {//
            ASSERT(e.enode == (uint32_t)(-1));
            if (sequence.Equal(e.start+active.len, seq_idx, wlen)) {
                active.len += wlen;
                return false;
            } else {
                return true;
            }
        } else {
            ASSERT(e.enode != (uint32_t)(-1));
            uint32_t elen = e.end - e.start + 1;
            ASSERT(active.len < elen);
            ASSERT(active.len + wlen <= elen);
            if (sequence.Equal(e.start + active.len, seq_idx, wlen)) {
                active.len += wlen;
                if (active.len == elen) {
                    active.node = e.enode;
                    active.edge = 0;
                    active.len = 0;
                }
                return false;
            } else {
                return true;
            }
        }
    }
}

template <typename SequenceElementType,
    typename SpaceParameter,
    typename SequenceClass,
    uint32_t WordLengthFunc(SequenceClass*, size_t),
    typename SequenceHashMethod>
bool ZSuffixTree<SequenceElementType, SpaceParameter, SequenceClass,
    WordLengthFunc, SequenceHashMethod>::Init()
{
    nodes.PushBack(ZSTNode<uint32_t>());
    active.node = 0;
    active.edge = 0;
    active.len = 0;
    return true;
}

#endif

